﻿using System;
using System.Diagnostics;
using System.Drawing;
using System.Runtime.InteropServices;
using WindowsMediaLib;
using WindowsMediaLib.Defs;

namespace Desktop.Media.Encoder
{
    [ComVisible(true)]
    public class WindowsMediaEncoder
    {
        #region Member variables
        private const short WAVE_FORMAT_PCM = 0x1;
        private IWMWriter m_wmriter;
        private IWMInputMediaProps m_videoInputProps;
        private IWMInputMediaProps m_audioInputProps;

        private int m_dwVideoInput = -1;
        private int m_dwAudioInput = -1;

        private string m_encodingProfileData;

        private long? m_streamBaseTime;
        private long? m_lastVideoTime;
        private long m_currentVideoTime;
        #endregion

        [DllImport("Kernel32.dll", EntryPoint = "RtlMoveMemory")]
        private static extern void CopyMemory(IntPtr destination, IntPtr source, [MarshalAs(UnmanagedType.U4)] int length);

        [DllImport("Kernel32.dll")]
        private static extern bool QueryPerformanceCounter(out long lpPerformanceCount);

        [DllImport("Kernel32.dll")]
        private static extern bool QueryPerformanceFrequency(out long lpFrequency);

        private static long m_timerFrequency;

        static WindowsMediaEncoder()
        {
            QueryPerformanceFrequency(out m_timerFrequency);
        }

        public int TcpPort { get; set; }

        public bool NetworkStream { get; set; }

        public string OutputFilename { get; set; }

        public bool UseGeneratedTimecodes { get; set; }

        public string EncodingProfileData
        {
            get { return m_encodingProfileData; }
            set
            {
                if (m_encodingProfileData == value) 
                    return;

                m_encodingProfileData = value;
                Initialize();
            }
        }

        private void Initialize()
        {
            WriteLog(string.Format("Initialize "));
            Close();
            Guid guidInputType;
            int dwInputCount;
            
            IWMProfileManager pWMProfileManager = null;
            IWMProfile pWMProfile = null;
            
            // Initialize all member variables
            m_dwVideoInput = -1;
            m_dwAudioInput = -1;
            
            m_wmriter = null;
            m_videoInputProps = null;
            m_audioInputProps = null;

            try
            {
                // Open the profile manager
                WMUtils.WMCreateProfileManager(out pWMProfileManager);

                // Convert pWMProfileManager to a IWMProfileManager2
                var pProfileManager2 = pWMProfileManager as IWMProfileManager2;

                // Specify the version number of the profiles to use
                pProfileManager2.SetSystemProfileVersion(WMVersion.V9_0);

                string profile = EncodingProfileData;

                pProfileManager2.LoadProfileByData(profile, out pWMProfile);

                // Create a writer.  This is the interface we actually write with
                WMUtils.WMCreateWriter(IntPtr.Zero, out m_wmriter);

                // Set the profile we got into the writer.  This controls compression, video
                // size, # of video channels, # of audio channels, etc
                m_wmriter.SetProfile(pWMProfile);
               
                // Find out how many inputs are in the current profile
                m_wmriter.GetInputCount(out dwInputCount);

                // Assume we won't find any video pins
                m_dwVideoInput = -1;
                m_dwAudioInput = -1;

                IWMInputMediaProps inputProps;

                // Find the first video input on the writer
                for (int i = 0; i < dwInputCount; i++)
                {
                    // Get the properties of channel #i
                    m_wmriter.GetInputProps(i, out inputProps);

                    // Read the type of the channel
                    inputProps.GetType(out guidInputType);

                    // If it is video, we are done
                    if (guidInputType == MediaType.Video)
                    {
                        m_videoInputProps = inputProps;
                        m_dwVideoInput = i;
                    }
                    else if(guidInputType == MediaType.Audio)
                    {
                        m_audioInputProps = inputProps;
                        m_dwAudioInput = i;
                    }
                }

                // Didn't find a channel
                if (m_videoInputProps == null && m_audioInputProps == null)
                    throw new Exception("Profile does not accept any input");

                var writerAdvanced2 = m_wmriter as IWMWriterAdvanced2;

                if (NetworkStream && writerAdvanced2 != null)
                {
                    IWMWriterNetworkSink netSink;
                    WMUtils.WMCreateWriterNetworkSink(out netSink);

                    netSink.SetNetworkProtocol(NetProtocol.HTTP);

                    int port = TcpPort;
                    netSink.Open(ref port);

                    writerAdvanced2.AddSink(netSink);
                }

                if (!string.IsNullOrEmpty(OutputFilename))
                {
                    m_wmriter.SetOutputFilename(OutputFilename);
                }
            }
            catch
            {
                Close();
                throw;
            }
            finally
            {
                if (pWMProfile != null)
                    Marshal.ReleaseComObject(pWMProfile);
                if (pWMProfileManager != null)
                    Marshal.ReleaseComObject(pWMProfileManager);
            }
        }

        public WindowsMediaEncoder()
        {
            TcpPort = 8080;
            NetworkStream = false;
        }

        /// <summary>
        /// Destructor
        /// </summary>
        ~WindowsMediaEncoder()
        {
            Close();
        }

        /// <summary>
        /// Close the output and release the variables
        /// </summary>
        public void Close()
        {
            m_streamBaseTime = null;
            m_lastVideoTime = null;
            m_currentVideoTime = 0;

            WriteLog(string.Format("Close buffer m_wmriter null: {0}", m_wmriter == null));
            if (m_videoInputProps != null)
            {
                Marshal.ReleaseComObject(m_videoInputProps);
                m_videoInputProps = null;
            }

            if (m_wmriter != null)
            {
                m_wmriter.Flush();
                m_wmriter.EndWriting();
                Marshal.ReleaseComObject(m_wmriter);
                m_wmriter = null;
            }
        }

        public void AppendAudioSample(byte[] sample, long sampleTime)
        {
            if (m_dwAudioInput == -1)
                return;

            INSSBuffer pSample = null;

            try
            {
                int iSize = sample.Length;
                IntPtr ip;

                m_wmriter.AllocateSample(iSize, out pSample);

                pSample.GetBuffer(out ip);

                Marshal.Copy(sample, 0, ip, sample.Length);

                WriteSample(m_dwAudioInput, sampleTime, SampleFlag.CleanPoint, pSample);
            }
            catch (Exception ex)
            {

            }
            finally
            {
                if (pSample != null)
                    Marshal.ReleaseComObject(pSample);
            }
        }

        private void WriteLog(string message)
        {
            System.Diagnostics.Debug.WriteLine(message);
        }

        public void AppendVideoSampleBuffer(byte[] pixels, long sampleTime)
        {
            if (m_dwVideoInput == -1)
                return;

            INSSBuffer pSample = null;

            try
            {
                int pBufferLength = pixels.Length;

                IntPtr ip;

                // Get a sample interface
                m_wmriter.AllocateSample(pBufferLength, out pSample);

                // Get the buffer from the sample interface.  This is
                // where we copy the bitmap data to
                pSample.GetBuffer(out ip);

                pSample.SetLength(pBufferLength);

                // Copy the bitmap data into the sample buffer
                Marshal.Copy(pixels, 0, ip, pixels.Length);

                WriteSample(m_dwVideoInput, sampleTime, SampleFlag.CleanPoint, pSample);
            }
            catch (Exception ex)
            {

            }
            finally
            {
                if (pSample != null)
                    Marshal.ReleaseComObject(pSample);
            }
        }

        public void AppendVideoSample(int[] pixels, long sampleTime)
        {
            if (m_dwVideoInput == -1)
                return;

            INSSBuffer pSample = null;
          
            try
            {
                int iSize = pixels.Length * sizeof(int);

                IntPtr ip;

                // Get a sample interface
                m_wmriter.AllocateSample(iSize, out pSample);

                // Get the buffer from the sample interface.  This is
                // where we copy the bitmap data to
                pSample.GetBuffer(out ip);

                // Copy the bitmap data into the sample buffer
                Marshal.Copy(pixels, 0, ip, pixels.Length);

                WriteSample(m_dwVideoInput, sampleTime, SampleFlag.CleanPoint, pSample);
            }
            catch(Exception ex)
            {
                
            }
            finally
            {
                if (pSample != null)
                    Marshal.ReleaseComObject(pSample);
            }
        }

        private void WriteSample(int input, long sampleTime, SampleFlag flags, INSSBuffer pSample)
        {
            if (m_streamBaseTime == null)
            {
                long time;
                QueryPerformanceCounter(out time);
                m_streamBaseTime = time;
            }

            /* Generate our own time code because apparent sl 
             * quirk giving us 0 on all presentation times on some web cams */
            if (input == m_dwVideoInput)
            {
                long currentTime;
                QueryPerformanceCounter(out currentTime);

                if (m_lastVideoTime == null)
                    m_lastVideoTime = currentTime;

                long timeDiff = (currentTime - m_lastVideoTime.Value);
                m_lastVideoTime = currentTime;

                double elapsedSeconds = (double)timeDiff / m_timerFrequency;

                m_currentVideoTime += (long) (elapsedSeconds*10000000);
                
                if(UseGeneratedTimecodes)
                    sampleTime = m_currentVideoTime;
            }

            m_wmriter.WriteSample(input, sampleTime, flags, pSample);
        }

        /// <summary>
        /// Create a WaveFormatEx from parameters
        /// </summary>
        /// <param name="channels">1 = mono, 2 = stereo</param>
        /// <param name="bits">8 = 8bit audio, 16 = 16bit audio</param>
        /// <param name="sampleRate">Audio rate (11025, 22050, 44100, 48000 etc)</param>
        private static WaveFormatEx CreateWaveFormatEx(short channels, short bits, int sampleRate)
        {
            var waveformat = new WaveFormatEx();
            waveformat.wFormatTag = WAVE_FORMAT_PCM;
            waveformat.nChannels = channels;
            waveformat.nSamplesPerSec = sampleRate;
            waveformat.wBitsPerSample = bits;
            waveformat.nBlockAlign = (short)(waveformat.nChannels * (waveformat.wBitsPerSample / 8));
            waveformat.nAvgBytesPerSec = waveformat.nSamplesPerSec * waveformat.nBlockAlign;
            waveformat.cbSize = 0;

            return waveformat;
        }

        public void InitializeVideo(int width, int height, int bitsPerPixel)
        {
            WriteLog(string.Format("InitializeVideo width: {0} height: {1} bpp {2} ", width, height, bitsPerPixel));
            if (m_dwVideoInput == -1)
                return;
          
            var mt = new AMMediaType();
            var videoInfo = new VideoInfoHeader();

            // Create the VideoInfoHeader using info from the bitmap
            videoInfo.BmiHeader.Size = Marshal.SizeOf(typeof(BitmapInfoHeader));
            videoInfo.BmiHeader.Width = width;
            videoInfo.BmiHeader.Height = height;
            videoInfo.BmiHeader.Planes = 1;

            // compression thru clrimportant don't seem to be used. Init them anyway
            videoInfo.BmiHeader.Compression = 0;
            videoInfo.BmiHeader.ImageSize = width * height * bitsPerPixel;
            videoInfo.BmiHeader.XPelsPerMeter = 0;
            videoInfo.BmiHeader.YPelsPerMeter = 0;
            videoInfo.BmiHeader.ClrUsed = 0;
            videoInfo.BmiHeader.ClrImportant = 0;
            videoInfo.BmiHeader.BitCount = (short)bitsPerPixel;

            switch (bitsPerPixel)
            {
                case 32:
                    mt.subType = MediaSubType.RGB32;
                    break;
                case 24:
                    mt.subType = MediaSubType.RGB24;
                    break;
                case 16:
                    mt.subType = MediaSubType.RGB555;
                    break;
                default:
                    throw new Exception("Unrecognized Pixelformat");
            }

            videoInfo.SrcRect = new Rectangle(0, 0, width, height);
            videoInfo.TargetRect = videoInfo.SrcRect;
            videoInfo.BmiHeader.ImageSize = width * height * (videoInfo.BmiHeader.BitCount / 8);
            videoInfo.BitRate = videoInfo.BmiHeader.ImageSize * 30;
            videoInfo.BitErrorRate = 0;
            videoInfo.AvgTimePerFrame = 10000000 / 30;

            mt.majorType = MediaType.Video;
            mt.fixedSizeSamples = true;
            mt.temporalCompression = false;
            mt.sampleSize = videoInfo.BmiHeader.ImageSize;
            mt.formatType = FormatType.VideoInfo;
            mt.unkPtr = IntPtr.Zero;
            mt.formatSize = Marshal.SizeOf(typeof(VideoInfoHeader));

            // Lock the videoInfo structure, and put the pointer
            // into the mediatype structure
            GCHandle gHan = GCHandle.Alloc(videoInfo, GCHandleType.Pinned);

            try
            {
                // Set the inputprops using the structures
                mt.formatPtr = gHan.AddrOfPinnedObject();

                m_videoInputProps.SetMediaType(mt);
            }
            catch (Exception ex)
            {
            }
            finally
            {
                gHan.Free();
                mt.formatPtr = IntPtr.Zero;
            }

            // Now take the inputprops, and set them on the file writer
            m_wmriter.SetInputProps(m_dwVideoInput, m_videoInputProps);
        }

        public void BeginWriting()
        {
            WriteLog(string.Format("BeginWriting"));
            m_wmriter.BeginWriting();
        }

        public void InitializeAudio(short channels, short bits, int sampleRate)
        {
            WriteLog(string.Format("InitializeVideo channels: {0} bits: {1} sampleRate: {2} ", channels, bits, sampleRate));

            if (m_dwAudioInput == -1)
                return;

            var mt = new AMMediaType();

            var waveformat = CreateWaveFormatEx(channels, bits, sampleRate);

            mt.majorType = MediaType.Audio;
            mt.subType = MediaSubType.PCM;
            mt.fixedSizeSamples = true;
            mt.temporalCompression = false;
            mt.sampleSize = (channels*bits*sampleRate)/8;
            mt.formatType = FormatType.WaveEx;
            mt.unkPtr = IntPtr.Zero;
            mt.formatSize = Marshal.SizeOf(typeof(WaveFormatEx));

            GCHandle gHan = GCHandle.Alloc(waveformat, GCHandleType.Pinned);

            try
            {
                // Set the inputprops using the structures
                mt.formatPtr = gHan.AddrOfPinnedObject();
                m_audioInputProps.SetMediaType(mt);
            }
            catch (Exception ex)
            {

            }
            finally
            {
                gHan.Free();
                mt.formatPtr = IntPtr.Zero;
            }

            m_wmriter.SetInputProps(m_dwAudioInput, m_audioInputProps);
        }
    }
}
